package edu.arizona.sista.learn.activity.fluents.paths;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Vector;

import edu.arizona.sista.learn.activity.model.Predicate;
import edu.arizona.sista.learn.activity.model.args.PredArg;
import edu.arizona.sista.learn.activity.util.PredicateNames;

public class PredicateExtractor2 extends PredicateExtractor {
	
	public PredicateExtractor2(Scene scene, DataType type) {
		super(scene, type);
	}
	
	public Vector<Vector<Predicate>> extractPredicates() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>(scene.getNumFrames());
		for (int i = 0; i < scene.getNumFrames(); i++) {
			retval.add(new Vector<Predicate>());
		}
		
		Vector<Vector<Predicate>> moving = extractMoving(0.03, 5);
		Vector<Vector<Predicate>> movingFastHuman = extractMovingFastHuman(0.235, 5);
		Vector<Vector<Predicate>> distances = extractDistances(0.01, 4);
		Vector<Vector<Predicate>> distancesToStationary = extractDistancesToStationary(0.01, 4);
		Vector<Vector<Predicate>> trajectoryStable = extractTrajectoryStable(0.09, 4);
		Vector<Vector<Predicate>> speeds = extractSpeeds(0.0005, 4);	// Probably still need more tweaking
		
		for (int i = 0; i < scene.getNumFrames(); i++) {
			retval.get(i).addAll(moving.get(i));
			retval.get(i).addAll(movingFastHuman.get(i));
			retval.get(i).addAll(distances.get(i));
			retval.get(i).addAll(distancesToStationary.get(i));
			retval.get(i).addAll(trajectoryStable.get(i));
			retval.get(i).addAll(speeds.get(i));
		}
		
		return retval;
	}
	
	public Vector<Vector<Predicate>> extractMoving(double threshold, int window) {
		Map<PredArg, boolean[]> moving = new HashMap<PredArg, boolean[]>();
		Map<PredArg, boolean[]> notMoving = new HashMap<PredArg, boolean[]>();
		
		for (SceneEntity e : scene.getEntities()) {
			boolean[] m = new boolean[scene.getNumFrames()];
			boolean[] nm = new boolean[scene.getNumFrames()];
			
			int ms = -1, me = -1, nms = -1, nme = -1, lost = -1;
			for (int i = 0; i <= scene.getNumFrames()-window; i++) {
				Point3D p1 = e.getPoint(i);
				Point3D p2 = e.getPoint(i+(window-1));
				
				if (p1 != null && p2 != null) {
					lost = -1;
					if (p1.distanceTo(p2) >= threshold) {
						// moving
						if (nms != -1) {
							assert (ms == -1) : "Can't be not moving & moving at same time!!";
							// end not-moving
							for (int j = nms; j < i; j++) {
								nm[j] = true;
							}
							nms = -1; nme = -1;
						}
						if (ms == -1) {
							ms = i;
						}
						me = i+(window-1);
					} else {
						// not-moving
						if (ms != -1) {
							// But was previously moving
							if (me < i) {
								// stopped moving
								for (int j = ms; j <= me; j++) {
									m[j] = true;
								}
								ms = -1; me = -1;
								
								// therefore, started not-moving
								assert (me+1 <= i) : "Something went wrong here!";
								nms = i;
							} else {
								// could potentially still be moving,
								// so we do nothing here
								continue;
							}
						} else if (nms == -1) {
							// start not-moving
							nms = i;
						}
						nme = i;
					}
					
				} else {
					// at least one of the points doesn't exist
					if (lost == -1)		// record when we first lost one of the points
						lost = i;
					if (nms != -1 && nme < i) {
						// stop not-moving if we lost enough points
						assert (ms == -1) : "Can't be not moving & moving at same time!!";
						assert (nms < lost) : "Can't start not moving if a point is lost!!";
						for (int j = nms; j < lost; j++) {
							nm[j] = true;
						}
						nms = -1; nme = -1;
					} else if (ms != -1 && me < i) {
						// stop moving if we lost enough points
						assert (ms < lost) : "Can't start moving if a point is lost!!";
						for (int j = ms; j < lost; j++) {
							m[j] = true;
						}
						ms = -1; me = -1;
					}
				}
			}
			
			// Clean up here
			if (nms != -1) {
				assert (ms == -1) : "Can't be not moving & moving at same time!!";
				for (int j = nms; j < scene.getNumFrames(); j++) {
					nm[j] = true;
				}
			} else if (ms != -1) {
				for (int j = ms; j < scene.getNumFrames(); j++) {
					m[j] = true;
				}
			}
			
			moving.put(e.getPredArg(), m);
			notMoving.put(e.getPredArg(), nm);
		}
		
		// Compile results
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		for (int i = 0; i < scene.getNumFrames(); i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			
			for (PredArg am : moving.keySet()) {
				if (moving.get(am)[i]) {
					frame.add(new Predicate(PredicateNames.MOVING, am));
				}
			}
			
			for (PredArg anm : notMoving.keySet()) {
				if (notMoving.get(anm)[i]) {
					frame.add(new Predicate(PredicateNames.NOT_MOVING, anm));
				}
			}
			
			retval.add(frame);
		}
		
		return retval;
	}
	
	/*
	 * Extract all the MOVING_FAST_HUMAN predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractMovingFastHuman(double threshold, int window) {
		Map<PredArg, boolean[]> moving = new HashMap<PredArg, boolean[]>();
		
		for (SceneEntity e : scene.getEntities()) {
			boolean[] m = new boolean[scene.getNumFrames()];
			
			int ms = -1, me = -1, lost = -1;
			for (int i = 0; i <= scene.getNumFrames()-window; i++) {
				Point3D p1 = e.getPoint(i);
				Point3D p2 = e.getPoint(i+(window-1));
				
				if (p1 != null && p2 != null) {
					lost = -1;
					if (p1.distanceTo(p2) >= threshold) {
						// moving
						if (ms == -1) {
							ms = i;
						}
						me = i+(window-1);
					} else {
						// not-moving
						if (ms != -1 && me < i) {
							// previously moving but now stopped
							for (int j = ms; j <= me; j++) {
								m[j] = true;
							}
							ms = -1; me = -1;
						}
					}
					
				} else {
					// at least one of the points doesn't exist
					if (lost == -1)		// record when we first lost one of the points
						lost = i;
					if (ms != -1 && me < i) {
						// stop moving if we lost enough points
						assert (ms < lost) : "Can't start moving if a point is lost!!";
						for (int j = ms; j < lost; j++) {
							m[j] = true;
						}
						ms = -1; me = -1;
					}
				}
			}
			
			// Clean up here
			if (ms != -1) {
				for (int j = ms; j < scene.getNumFrames(); j++) {
					m[j] = true;
				}
			}
			
			moving.put(e.getPredArg(), m);
		}
		
		// Compile results
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		for (int i = 0; i < scene.getNumFrames(); i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			
			for (PredArg am : moving.keySet()) {
				if (moving.get(am)[i]) {
					frame.add(new Predicate(PredicateNames.MOVING_FAST_HUMAN, am));
				}
			}
			retval.add(frame);
		}
		
		return retval;
	}

	/*
	 * Extract all the DISTANCE_INCREASING/DECREASING/STABLE predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractDistances(double threshold, int window) {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		// pre-pad vector
		for (int i = 0; i < window/2; i++) {
			retval.add(new Vector<Predicate>());
		}
		
		for (int i = window/2; i < scene.getNumFrames()-(window/2); i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity[] pair : pairPermutations()) {
				// Get average distances from before, and after
				int counter1 = 0;
				double d1 = -1.0;
				for (int j = i - (window/2); j < i; j++) {
					Point3D e1 = pair[0].getPoint(j);
					Point3D e2 = pair[1].getPoint(j);
					if (e1 != null && e2 != null) {
						counter1++;
						d1 += e1.distanceTo(e2);
					}
				}
				if (counter1 > 0)
					d1 = d1 / counter1;
				
				int counter2 = 0;
				double d2 = -1.0;
				for (int j = i+1; j <= i + (window/2); j++) {
					Point3D e1 = pair[0].getPoint(j);
					Point3D e2 = pair[1].getPoint(j);
					if (e1 != null && e2 != null) {
						counter2++;
						d2 += e1.distanceTo(e2);
					}
				}
				if (counter2 > 0)
					d2 = d2 / counter2;
				
				if (counter1 > 0 && counter2 > 0) {
					if (d2 - d1 >= threshold) {
						frame.add(new Predicate(PredicateNames.DISTANCE_INCREASING, pair[0].getPredArg(), pair[1].getPredArg()));
					} else if (d2 - d1 <= -threshold) {
						frame.add(new Predicate(PredicateNames.DISTANCE_DECREASING, pair[0].getPredArg(), pair[1].getPredArg()));
					} else {
						frame.add(new Predicate(PredicateNames.DISTANCE_STABLE, pair[0].getPredArg(), pair[1].getPredArg()));
					}
				}
			}
			retval.add(frame);
		}
		
		// pad vector to full
		for (int i = retval.size(); i < scene.getNumFrames(); i++) {
			retval.add(new Vector<Predicate>());
		}
		
		return retval;
	}
	
	/*
	 * Extract all the DISTANCE_INCREASING/DECREASING_TO_STATIONARY predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractDistancesToStationary(double threshold, int window) {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		// pre-pad vector
		for (int i = 0; i < window/2; i++) {
			retval.add(new Vector<Predicate>());
		}
		
		for (int i = window/2; i < scene.getNumFrames()-(window/2); i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity[] pair : pairPermutations()) {
				// Set one point as stationary
				Point3D e2 = pair[1].getPoint(i);

				// Get average distances to stationary from before, and after
				int counter1 = 0;
				double d1 = -1.0;
				for (int j = i - (window/2); j < i; j++) {
					Point3D e1 = pair[0].getPoint(j);
					if (e1 != null && e2 != null) {
						counter1++;
						d1 += e1.distanceTo(e2);
					}
				}
				if (counter1 > 0)
					d1 = d1 / counter1;
				
				int counter2 = 0;
				double d2 = -1.0;
				for (int j = i+1; j <= i + (window/2); j++) {
					Point3D e1 = pair[0].getPoint(j);
					if (e1 != null && e2 != null) {
						counter2++;
						d2 += e1.distanceTo(e2);
					}
				}
				if (counter2 > 0)
					d2 = d2 / counter2;
				
				if (counter1 > 0 && counter2 > 0) {
					if (d2 - d1 >= threshold) {
						frame.add(new Predicate(PredicateNames.DISTANCE_INCREASING_TO_STATIONARY, pair[0].getPredArg(), pair[1].getPredArg()));
					} else if (d2 - d1 <= -threshold) {
						frame.add(new Predicate(PredicateNames.DISTANCE_DECREASING_TO_STATIONARY, pair[0].getPredArg(), pair[1].getPredArg()));
					}
				}
			}
			retval.add(frame);
		}
		
		// pad vector to full
		for (int i = retval.size(); i < scene.getNumFrames(); i++) {
			retval.add(new Vector<Predicate>());
		}
		
		return retval;
	}
		
	/*
	 * Extract all the DISTANCE_INCREASING_TO_STATIONARY predicate occurrences in the scene.
	 */
	public Vector<Vector<Predicate>> extractDistanceIncreasingToStationary() {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		for (int i = 0; i < scene.getNumFrames()-1; i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity[] pair : pairPermutations()) {
				Point3D p1e1 = pair[0].getPoint(i);
				Point3D p1e2 = pair[1].getPoint(i);
				Point3D p2e1 = pair[0].getPoint(i+1);
				
				if (p1e1 == null || p1e2 == null || p2e1 == null) {
					continue;
				}
				
				double d1 = p1e1.distanceTo(p1e2);
				double d2 = p2e1.distanceTo(p1e2);
				
				if (d2 - d1 >= getThreshold(PredicateNames.DISTANCE_INCREASING_TO_STATIONARY, pair[0])) {
					frame.add(new Predicate(PredicateNames.DISTANCE_INCREASING_TO_STATIONARY, pair[0].getPredArg(), pair[1].getPredArg()));
				}
			}
			retval.add(frame);
		}
		retval.add(new Vector<Predicate>());
		
		return retval;
	}	
	
	
	public Vector<Vector<Predicate>> extractTrajectoryStable(double threshold, int window) {
		Map<PredArg, boolean[]> trajstable = new HashMap<PredArg, boolean[]>();
		
		for (SceneEntity e : scene.getEntities()) {
			boolean[] stable = new boolean[scene.getNumFrames()];
			
			int ts = -1, te = -1, lost = -1;
			for (int i = window/2; i < scene.getNumFrames()-(window/2); i++) {
				Point3D p1 = e.getPoint(i-(window/2));
				Point3D p2 = e.getPoint(i);
				Point3D p3 = e.getPoint(i+(window/2));
				
				if (p1 != null && p2 != null && p3 != null) {
					lost = -1;
					Vec3D v1 = new Vec3D(p1, p2);
					Vec3D v2 = new Vec3D(p2, p3);
					double angle = v1.angle(v2);
					if ((Math.abs(angle)) < threshold) {
						// stable
						if (ts == -1) {
							ts = i;
						}
						te = i;
					} else {
						// not-stable
						if (ts != -1 && te < i) {
							// previously moving but now stopped
							for (int j = ts; j <= te; j++) {
								stable[j] = true;
							}
							ts = -1; te = -1;
						}
					}
					
				} else {
					// at least one of the points doesn't exist
					if (lost == -1)		// record when we first lost one of the points
						lost = i;
					if (ts != -1 && te < i) {
						// stop moving if we lost enough points
						assert (ts < lost) : "Can't start moving if a point is lost!!";
						for (int j = ts; j < lost; j++) {
							stable[j] = true;
						}
						ts = -1; te = -1;
					}
				}
			}
			
			// Clean up here
			if (ts != -1) {
				for (int j = ts; j < scene.getNumFrames(); j++) {
					stable[j] = true;
				}
			}
			
			trajstable.put(e.getPredArg(), stable);
		}
		
		// Compile results
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		for (int i = 0; i < scene.getNumFrames(); i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			
			for (PredArg a : trajstable.keySet()) {
				if (trajstable.get(a)[i]) {
					frame.add(new Predicate(PredicateNames.TRAJECTORY_STABLE, a));
				}
			}
			retval.add(frame);
		}
		
		return retval;
	}
	
	public Vector<Vector<Predicate>> extractSpeeds(double threshold, int window) {
		Vector<Vector<Predicate>> retval = new Vector<Vector<Predicate>>();
		
		// pre-pad vector
		for (int i = 0; i < window/2; i++) {
			retval.add(new Vector<Predicate>());
		}
		
		for (int i = window/2; i < scene.getNumFrames()-(window/2); i++) {
			Vector<Predicate> frame = new Vector<Predicate>();
			for (SceneEntity e : scene.getEntities()) {
				Point3D p2 = e.getPoint(i);
				
				// Get average distance from before to current, and current to after
				int counter1 = 0;
				double d1 = -1.0;
				for (int j = i - (window/2); j < i; j++) {
					Point3D p1 = e.getPoint(j);
					if (p1 != null && p2 != null) {
						counter1++;
						d1 += p1.distanceTo(p2);
					}
				}
				if (counter1 > 0)
					d1 = d1 / counter1;
				
				int counter2 = 0;
				double d2 = -1.0;
				for (int j = i+1; j <= i + (window/2); j++) {
					Point3D p3 = e.getPoint(j);
					if (p2 != null && p3 != null) {
						counter2++;
						d2 += p2.distanceTo(p3);
					}
				}
				if (counter2 > 0)
					d2 = d2 / counter2;
				
				if (counter1 > 0 && counter2 > 0) {
					if (d2 - d1 >= threshold) {
						frame.add(new Predicate(PredicateNames.SPEED_INCREASING, e.getPredArg()));
					} else if (d2 - d1 <= -threshold) {
						frame.add(new Predicate(PredicateNames.SPEED_DECREASING, e.getPredArg()));
					}
				}
			}
			retval.add(frame);
		}
		
		// pad vector to full
		for (int i = retval.size(); i < scene.getNumFrames(); i++) {
			retval.add(new Vector<Predicate>());
		}
		
		return retval;
	}
	
	// returns a Vector of pairs of SceneEntity
	// it can be used for the detection of predicates like distance-increasing
	// in which we have to compare a SceneEntity to every other one
	private Vector<SceneEntity[]> pairPermutations() {
		Vector<SceneEntity[]> pairs = new Vector<SceneEntity[]>();
		Iterator<SceneEntity> i1 = scene.getEntities().iterator();
		while (i1.hasNext()) {
			SceneEntity e1 = i1.next();
			Iterator<SceneEntity> i2 = scene.getEntities().iterator();
			while (i2.hasNext()) {
				SceneEntity e2 = i2.next();
				if (!e1.equals(e2)) {
					SceneEntity[] pair = new SceneEntity[2];
					pair[0] = e1;
					pair[1] = e2;
					pairs.add(pair);
				}
			}
		}
		return pairs;
	}
}
